import argparse

import numpy as np
from osgeo.gdal import Open, GetDriverByName, GDT_Int32
from typing import List


def majority_filter(center: int, neighbors: List[int], center_type="body"):
    """
    四邻域众数滤波。

    :param center: 中心像元的值。
    :param neighbors: 邻域像元的值，以任意次序排列均可。
    :param center_type: 中心像元的类型，可以是 'body','edge','corner' 中的一种。
    :return: 滤波后中心像元的值。
    """
    if center_type == 'corner':
        if neighbors[0] == neighbors[1]:
            return neighbors[0]
        else:
            return center
    else:
        if center_type == 'edge':
            th = 2
        else:
            th = 3
        for k in np.unique(neighbors):
            count_k = neighbors.count(k)
            if count_k >= th:
                # k占据多数，中心像素值替换为k
                return k
        return center


class MajorityFilter:
    def __init__(self, args, **kwargs):
        self.input_path = None
        self.output_path = None
        if args is not None:
            for key in args.__dict__:
                self[key] = args.__dict__[key]
        for key in kwargs.keys():
            self[key] = kwargs[key]

    def __getitem__(self, item):
        return self.__dict__[item]

    def __setitem__(self, key, value):
        self.__dict__[key] = value

    def __str__(self):
        return f"input={self.input_path},output={self.output_path}"

    def run(self):
        """
        执行滤波操作。

        :return:
        """
        # 读取数据
        img_src = Open(self.input_path)
        proj = img_src.GetProjection()
        trans = img_src.GetGeoTransform()
        n_rows, n_cols = img_src.RasterYSize, img_src.RasterXSize
        nodata = img_src.GetRasterBand(1).GetNoDataValue()
        img = img_src.ReadAsArray()
        filtered = np.copy(img)
        # 分角、边、主体区域进行滤波
        for r in range(n_rows):
            for c in range(n_cols):
                if r == 0:
                    if c == 0:
                        # 左上角
                        ct = 'corner'
                        nbs = [img[r, c + 1], img[r + 1, c]]
                    elif c == n_cols - 1:
                        # 右上角
                        ct = 'corner'
                        nbs = [img[r, c - 1], img[r + 1, c]]
                    else:
                        # 上边缘
                        ct = 'edge'
                        nbs = [img[r, c - 1], img[r, c + 1], img[r + 1, c]]
                elif r == n_rows - 1:
                    if c == 0:
                        # 左下角
                        ct = 'corner'
                        nbs = [img[r - 1, c], img[r, c + 1]]
                    elif c == n_cols - 1:
                        # 右下角
                        ct = 'corner'
                        nbs = [img[r - 1, c], img[r, c - 1]]
                    else:
                        # 下边缘
                        ct = 'edge'
                        nbs = [img[r, c - 1], img[r, c + 1], img[r - 1, c]]
                else:
                    if c == 0:
                        # 左边缘
                        ct = 'edge'
                        nbs = [img[r - 1, c], img[r + 1, c], img[r, c + 1]]
                    elif c == n_cols - 1:
                        # 右边缘
                        ct = 'edge'
                        nbs = [img[r - 1, c], img[r + 1, c], img[r, c - 1]]
                    else:
                        # body
                        ct = 'body'
                        nbs = [img[r - 1, c], img[r + 1, c], img[r, c - 1], img[r, c + 1]]
                filtered[r, c] = majority_filter(img[r, c], neighbors=nbs, center_type=ct)
        # 输出滤波结果
        driver = GetDriverByName("GTiff")
        img_out = driver.Create(self.output_path, n_cols, n_rows, 1, GDT_Int32)
        img_out.SetGeoTransform(trans)
        img_out.SetProjection(proj)
        img_out.GetRasterBand(1).WriteArray(filtered)
        img_out.GetRasterBand(1).SetNoDataValue(nodata)
        del img_out


if __name__ == '__main__':
    # 接收控制台传入参数
    parser = argparse.ArgumentParser(description='classification parameters')
    parser.add_argument('input_path', type=str, help='input file path')
    parser.add_argument('output_path', type=str, help='output file path')
    args = parser.parse_args()
    # print(args.__dict__)
    mf = MajorityFilter(args=args)
    mf.run()
    # print(slp)
